{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "\n",
    "import torch.distributed as dist\n",
    "import torch.optim as optim\n",
    "import torch.optim.lr_scheduler as lr_scheduler\n",
    "\n",
    "import test  # import test.py to get mAP after each epoch\n",
    "from models import *\n",
    "from utils.datasets import *\n",
    "from utils.utils import *\n",
    "\n",
    "from mymodel import *\n",
    "\n",
    "# Hyperparameters (results68: 59.9 mAP@0.5 yolov3-spp-416) https://github.com/ultralytics/yolov3/issues/310\n",
    "\n",
    "hyp = {'giou': 3.54,  # giou loss gain\n",
    "       'cls': 37.4,  # cls loss gain\n",
    "       'cls_pw': 1.0,  # cls BCELoss positive_weight\n",
    "       'obj': 64.3,  # obj loss gain (*=img_size/320 if img_size != 320)\n",
    "       'obj_pw': 1.0,  # obj BCELoss positive_weight\n",
    "       'iou_t': 0.225,  # iou training threshold\n",
    "       'lr0': 0.01,  # initial learning rate (SGD=5E-3, Adam=5E-4)\n",
    "       'lrf': -4.,  # final LambdaLR learning rate = lr0 * (10 ** lrf)\n",
    "       'momentum': 0.937,  # SGD momentum\n",
    "       'weight_decay': 0.000484,  # optimizer weight decay\n",
    "       'fl_gamma': 0.5,  # focal loss gamma\n",
    "       'hsv_h': 0.0138,  # image HSV-Hue augmentation (fraction)\n",
    "       'hsv_s': 0.678,  # image HSV-Saturation augmentation (fraction)\n",
    "       'hsv_v': 0.36,  # image HSV-Value augmentation (fraction)\n",
    "       'degrees': 1.98,  # image rotation (+/- deg)\n",
    "       'translate': 0.05,  # image translation (+/- fraction)\n",
    "       'scale': 0.05,  # image scale (+/- gain)\n",
    "       'shear': 0.641}  # image shear (+/- deg)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "\n",
    "parser.add_argument('--batch-size', type=int, default=16)  # effective bs = batch_size * accumulate = 16 * 4 = 64\n",
    "parser.add_argument('--accumulate', type=int, default=4, help='batches to accumulate before optimizing')\n",
    "parser.add_argument('--cfg', type=str, default='cfg/yolov3-tiny-1cls_1.cfg', help='*.cfg path')\n",
    "parser.add_argument('--data', type=str, default='data/coco2017.data', help='*.data path')\n",
    "parser.add_argument('--img-size', nargs='+', type=int, default=[320], help='train and test image-sizes')\n",
    "parser.add_argument('--rect', action='store_true', help='rectangular training')\n",
    "parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')\n",
    "parser.add_argument('--weights', type=str, default='/home/denggc/DAC2021/dgc/April/ultra_bypass/weights/test_best.pt', help='initial weights path')\n",
    "parser.add_argument('--arc', type=str, default='default', help='yolo architecture')  # default, uCE, uBCE\n",
    "parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied')\n",
    "parser.add_argument('--device', default='1', help='device id (i.e. 0 or 0,1 or cpu)')\n",
    "parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')\n",
    "parser.add_argument('--var', type=float, help='debug variable')\n",
    "opt = parser.parse_known_args()[0]\n",
    "print(opt)\n",
    "print(opt.weights)\n",
    "device = torch_utils.select_device(opt.device, batch_size=opt.batch_size)\n",
    "print(device)\n",
    "\n",
    "\n",
    "img_size, img_size_test = opt.img_size if len(opt.img_size) == 2 else opt.img_size * 2  # train, test sizes\n",
    "batch_size = opt.batch_size\n",
    "accumulate = opt.accumulate  # effective bs = batch_size * accumulate = 16 * 4 = 64\n",
    "weights = opt.weights  # initial training weights\n",
    "\n",
    "\n",
    "test_path = '../DAC-SDC2021/dataset/sample'\n",
    "nc = 1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = UltraNet_Bypass().to(device)\n",
    "model.hyp = hyp\n",
    "model.nc = 1\n",
    "model.arc = 'default'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if weights.endswith('.pt'):  # pytorch format\n",
    "    # possible weights are '*.pt', 'yolov3-spp.pt', 'yolov3-tiny.pt' etc.\n",
    "    print(\"load weights...\")\n",
    "    model.load_state_dict(torch.load(weights, map_location=device)['model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = min(batch_size, 1)\n",
    "nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers\n",
    "dataset = LoadImagesAndLabels(test_path, img_size_test, batch_size * 2,\n",
    "                              hyp=hyp,\n",
    "                              rect=False,\n",
    "                              cache_images=opt.cache_images,\n",
    "                              single_cls=opt.single_cls)\n",
    "testloader = torch.utils.data.DataLoader(dataset,\n",
    "                                         batch_size=batch_size * 2,\n",
    "                                         num_workers=nw,\n",
    "                                         pin_memory=True,\n",
    "                                         collate_fn=dataset.collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = test.test(opt.cfg,\n",
    "                    opt.data,\n",
    "                    batch_size=batch_size * 2,\n",
    "                    img_size=img_size_test,\n",
    "                    model=model,\n",
    "                    conf_thres=0.001,  # 0.001 if opt.evolve or (final_epoch and is_coco) else 0.01,\n",
    "                    iou_thres=0.6,\n",
    "                    save_json=False,\n",
    "                    single_cls=opt.single_cls,\n",
    "                    dataloader=testloader)\n",
    "print(results)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
